#!/bin/sh

set -ex

# Setups
exp_id="exp1" 		# Make sure to name each experiment
n_iter=5 			# Adjust number of PTL iterations
device_ids=0,1,2,3	# Adjust gpu_ids
batch_size_gan=8	# Adjust bs for training CycleGAN to fit your environment
n_epoch_gan=100
n_epoch_gan_less=0  # New: Progressively train CycleGAN if != 0
n_iter_det_less=0   # New: Progressively train RetinaNet if != 0
step_det_less=0

# Dataset Initilization
# Make sure to produce the correct dataset e.g., --real_dataset_spl_num
python DatasetInit.py --exp_id ${exp_id} --real_dataset_spl_num 50

# PTL iterations
for ((idx=0; idx<n_iter; idx++)); do
	cur_iter="iter_${idx}"

	# Transformation Candidate Selection
	# Make sure to adjust some parameters for your environmnet e.g., --n_gpus, --batch_size_det
	if [ ${idx} -ne 0 ] && [ ${n_iter_det_less} -ne 0 ]; then
		python TransCandiSelect.py --exp_id ${exp_id} --cur_ptl_iter ${cur_iter} --n_iter_det ${n_iter_det_less} --step_det ${step_det_less} --progressive
	else	
		python TransCandiSelect.py --exp_id ${exp_id} --cur_ptl_iter ${cur_iter}
	fi

	# Virtual2Real Transformation
	if [ ${idx} -ne 0 ] && [ ${n_epoch_gan_less} -ne 0 ]; then
		mkdir -p ./models/${exp_id}/gan/${cur_iter}
		cp ./models/${exp_id}/gan/"iter_$((idx-1))"/latest_*.pth ./models/${exp_id}/gan/${cur_iter}/
		python ./pytorch-CycleGAN-and-pix2pix/train_gan.py --dataroot ./datasets/${exp_id}/gan/${cur_iter} --checkpoints_dir ./models/${exp_id}/gan --name ${cur_iter} --gpu_ids ${device_ids} --model cycle_gan --save_epoch_freq 20 \
		--netG resnet_9blocks --no_dropout --load_size 256 --batch_size ${batch_size_gan} --preprocess none --n_epochs ${n_epoch_gan_less} --n_epochs_decay 0 --pool_size 50 --display_ncols 0 --lr 0.0002 --display_id 0 --continue_train		
	else
		python ./pytorch-CycleGAN-and-pix2pix/train_gan.py --dataroot ./datasets/${exp_id}/gan/${cur_iter} --checkpoints_dir ./models/${exp_id}/gan --name ${cur_iter} --gpu_ids ${device_ids} --model cycle_gan --save_epoch_freq 20 \
		--netG resnet_9blocks --no_dropout --load_size 256 --batch_size ${batch_size_gan} --preprocess none --n_epochs ${n_epoch_gan} --n_epochs_decay 0 --pool_size 50 --display_ncols 0 --lr 0.0002 --display_id 0
	fi
	python ./pytorch-CycleGAN-and-pix2pix/test_gan.py --dataroot ./datasets/${exp_id}/gan/${cur_iter} --checkpoints_dir ./models/${exp_id}/gan --name ${cur_iter} --results_dir ./results/${exp_id}/gan --model cycle_gan \
	--netG resnet_9blocks --phase test --no_dropout --load_size 256 --preprocess none --eval --gpu_ids 0

	# Set Update
	python SetUpdate.py --exp_id ${exp_id} --cur_ptl_iter ${cur_iter}
done

# Train and evaluate the last detector
last_iter="iter_${n_iter}"
if [ ${idx} -ne 0 ] && [ ${n_iter_det_less} -ne 0 ]; then
	python TransCandiSelect.py --exp_id ${exp_id} --cur_ptl_iter ${last_iter} --n_iter_det ${n_iter_det_less} --step_det ${step_det_less} --progressive --det_only
else	
	python TransCandiSelect.py --exp_id ${exp_id} --cur_ptl_iter ${last_iter} --det_only
fi

# Keep the run file
mkdir -p ./runs/${exp_id}
cp "$0" ./runs/${exp_id}/

# After finishing one experiment, you can go ./utils/ and do the following to collect results
# python collect_ptl_results --exp_id ${exp_id} --n_ptl_iters n_iter -> Results will be generated at: ./results/exp1/det/results.txt